import glob
import os.path

import cv2
import numpy as np
import xml.etree.ElementTree as ET

import tqdm


def resize_and_paste_image(scale_img_path, xml_path, output_path, target_size=(416, 416)):
    # 创建纯黑色背景图
    background = np.zeros((target_size[1], target_size[0], 3), dtype=np.uint8)
    # background[:] = background_color
    # 读取前景图
    foreground = cv2.imread(scale_img_path)
    # 获取前景图原始大小
    original_height, original_width = foreground.shape[:2]
    # 计算缩放比例，保持原始纵横比
    scale_x = target_size[0] / original_width
    scale_y = target_size[1] / original_height
    scale = min(scale_x, scale_y)
    # 缩放前景图
    resized_foreground = cv2.resize(foreground, (int(original_width * scale), int(original_height * scale)))
    # 读取XML文件
    tree = ET.parse(xml_path)
    root = tree.getroot()
    # 更新图像大小
    size = root.find('size')
    w = size.find('width').text
    h = size.find('height').text
    w, h = int(w), int(h)
    if w == target_size[0] and h == target_size[1]:
        print(f"{xml_path}尺寸是{(w, h)}")
        return
    size.find('width').text = str(target_size[0])
    size.find('height').text = str(target_size[1])

    # 更新每个边界框的坐标
    for obj in root.findall('object'):
        bbox = obj.find('bndbox')
        xmin = int(bbox.find('xmin').text)
        ymin = int(bbox.find('ymin').text)
        xmax = int(bbox.find('xmax').text)
        ymax = int(bbox.find('ymax').text)

        # 缩放边界框坐标
        bbox.find('xmin').text = str(int((xmin - original_width // 2) * scale + target_size[0] // 2))
        bbox.find('ymin').text = str(int((ymin - original_height // 2) * scale + target_size[1] // 2))
        bbox.find('xmax').text = str(int((xmax - original_width // 2) * scale + target_size[0] // 2))
        bbox.find('ymax').text = str(int((ymax - original_height // 2) * scale + target_size[1] // 2))
    #     cv2.rectangle(foreground, (xmin, ymin), (xmax, ymax), color=(0, 0, 255), thickness=2)
    # cv2.imshow("foreground", foreground)
    # cv2.waitKey(0)
    # cv2.destroyAllWindows()

    # 获取贴图的位置
    x_offset = (target_size[0] - resized_foreground.shape[1]) // 2
    y_offset = (target_size[1] - resized_foreground.shape[0]) // 2
    # 将缩放后的前景图贴在背景图上
    background[y_offset:y_offset + resized_foreground.shape[0],
    x_offset:x_offset + resized_foreground.shape[1]] = resized_foreground
    # 保存合成后的图像
    cv2.imwrite(output_path, background)
    # 保存更新后的XML文件
    tree.write(xml_path)

def convert_img(foreground_path_dirs, xml_path_dirs, output_path):
    foreground_paths = glob.glob(os.path.join(foreground_path_dirs, "*"))
    xml_paths = glob.glob(os.path.join(xml_path_dirs, "*"))
    for idx, xml_path in tqdm.tqdm(enumerate(xml_paths), total=len(xml_paths)):
        foreground_path = foreground_paths[idx]
        img_name = foreground_path.split("\\")[-1]
        foreground_output_path = os.path.join(output_path, img_name)
        resize_and_paste_image(foreground_path, xml_path, foreground_output_path)
        pass


def show_img(xml_path, img_path):
    tree = ET.parse(xml_path)
    root = tree.getroot()
    bg_img = cv2.imread(img_path)
    # 更新每个边界框的坐标
    for obj in root.findall('object'):
        bbox = obj.find('bndbox')
        xmin = int(bbox.find('xmin').text)
        ymin = int(bbox.find('ymin').text)
        xmax = int(bbox.find('xmax').text)
        ymax = int(bbox.find('ymax').text)
        cv2.rectangle(bg_img, (xmin, ymin), (xmax, ymax), color=(0, 0, 255), thickness=2)
    if bg_img is None:
        print(f"{img_path} 不存在")
        return
    cv2.imshow("bg_img", bg_img)
    cv2.waitKey(0)
    cv2.destroyAllWindows()


def batch_scale_img(ori_image_dir, xml_path, output_path_dir):
    # 转换图片 xml_path 是原图片的voc(xml)标签目录
    # output_path是转换之后输出图片路径
    # ori_image_path 是原图像目录
    ori_image_paths = glob.glob(os.path.join(ori_image_dir, "*"))
    xml_paths = glob.glob(os.path.join(xml_path, "*"))
    zip_paths = zip(ori_image_paths, xml_paths)
    for ori_image_path, xml_path in tqdm.tqdm(zip_paths, total=len(xml_paths)):
        img_name = ori_image_path.split("\\")[-1]
        scale_output_path = os.path.join(output_path_dir, img_name)
        resize_and_paste_image(ori_image_path, xml_path, scale_output_path)


def scale_single_scale_img(ori_image_path, xml_path, output_path):
    # 单张图片进行缩放
    # ori_image_path = 'test/111.jpg'
    # xml_path = 'test/111.xml'
    # output_path = 'test/output_image_111.jpg'
    # print(f"修改前 {xml_path}")
    # 调用函数
    resize_and_paste_image(ori_image_path, xml_path, output_path)
    print(f"修改后 {xml_path}")
    show_img(xml_path, output_path)

if __name__ == '__main__':
    # 单张图片进行缩放并显示
    # ori_image_path = 'test/ori_4.jpg'
    # xml_path = 'test/4.xml'
    # output_path = 'test/output_image_4.jpg'
    # scale_single_scale_img(ori_image_path, xml_path, output_path)
    # 原图像地址
    ori_image_path = r"../data/VOC2007/JPEGImages"
    # 拷贝Annotations_bak目录下的标签到Annotations, 标签转换之后会被修改满足YOLO_v3标签信息 416x416, 然后xmin, ymin, xmax, ymax会进行等比例缩放
    xml_annotion_path = r"../data/VOC2007/Annotations"
    # 转换为416 * 416的图片存储缩放之后的目录
    scale_output_path_dir = r"../data/VOC2007/YOLOv3_JPEGImages"
    # 转换图片 xml_path 是原图片的voc(xml)标签路径
    # output_path是转换之后输出图片路径
    # ori_image_path 是原图像
    batch_scale_img(ori_image_path, xml_annotion_path, scale_output_path_dir)
    scale_img_paths = glob.glob(os.path.join(scale_output_path_dir, "*"))
    for idx, img_path in enumerate(scale_img_paths):
        name = img_path.split("\\")[-1].split(".")[0]
        xml_path = os.path.join(xml_annotion_path, f"{name}.xml")
        scale_img_path = scale_img_paths[idx]
        show_img(xml_path, scale_img_path)
